from flask import Blueprint, render_template, request, redirect, url_for, send_file, session, abort, flash
from flask_login import login_user, logout_user, current_user, login_required
from webapp.models import User
from io import BytesIO


user_blueprint = Blueprint(
    'user',
    __name__
)


@user_blueprint.route('/')
@login_required
def home():
    return render_template('user.html')


@user_blueprint.route('/login', methods=['POST', 'GET'])
def login():

    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']

        user = User.query.filter_by(username=username).first()
        if user.check_password(password):
            if user.totp_key:
                # 进行二次验证
                session['login_user_2step'] = user.id
                return redirect(url_for('user.two_step', next=request.args.get('next')))

            login_user(user)
            if request.args.get('next'):
                return redirect(request.args.get('next'))
            return redirect(url_for('pic.home'))
        else:
            pass

    return render_template('login.html')


@user_blueprint.route('/init')
def init_user():
    if not User.query.all():
        password = '123456'
        user = User(
            username='superyou',
            password=password,
            totp_key='',
        )
        return user.username, password
    return 'User already exists.'


@user_blueprint.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('user.login'))


@user_blueprint.route("/qr")
@login_required
def qrcode_img():
    user = current_user
    img = BytesIO()
    q = user.qrcode()
    q.save(img)
    img.seek(0)
    return send_file(img, mimetype="image/png")


@user_blueprint.route("/two_step", methods=['POST', 'GET'])
def two_step():
    if not session.get('login_user_2step'):
        if current_user:
            return redirect(url_for('pic.home'))
        else:
            return redirect(url_for('user.login'))
    
    if request.method == 'POST':
        user_id = session.get('login_user_2step', '')
        user = User.query.get_or_404(user_id)
        totp = request.form.get('totp_input')

        if not user.check_totp(totp):
            retry = session.get('2step_retry', 0)
            retry += 1
            session['2step_retry'] = retry
            if retry < 3:
                flash('code error', category='warning')
                return render_template('two_step.html')
            else:
                session.pop('2step_retry', '')
                flash('code error 3rd, please try login again', category='danger')
                return redirect(url_for('user.login'))

        else:
            session.pop('login_user_2step', '')
            session.pop('2step_retry', '')
            login_user(user)
            if request.args.get('next'):
                return redirect(request.args.get('next'))
            return redirect(url_for('pic.home'))

    return render_template('two_step.html')
